# this file plan a traj for lateral control within one cycle time based on sample alg
# the alg processing is as follows:
# 1. sample longitude velocity and lateral position
# 2. merge longitude and lateral traj
# 3. select optimize traj
# (1) filter invalid acceleration and yawrate
# (2) collision detecte
# (3) calc cost function return the minimun one

import matplotlib.pyplot as plt
import numpy as np
from sympy import *
import scipy.integrate as intg
import math
import time

# struct to sort coeff of a traj about sample time and lateral and longitude
class coeff_strut:
    def __init__(self,lateral_coeff,longitude_coeff,pred_t):
        self.lateral_coeff=lateral_coeff
        self.longitude_coeff=longitude_coeff
        self.predict_t=pred_t
    def get_lateral_coeff(self):
        return self.lateral_coeff
    def get_longitude_coeff(self):
        return self.longitude_coeff
    def get_predt(self):
        return self.predict_t

class TrajCoeff:
    def __init__(self,coeff,sampletime):
        self.coeff=coeff
        self.sampletime=sampletime
    def getCoeff(self):
        return self.coeff
    def getSampletime(self):
        return self.sampletime
# struct to store vehicle kinetics state
class vehicle_state:
    def __init__(self,x,y,vx,vy,ax,ay):
        self.x=x
        self.y=y
        self.vx=vx
        self.vy=vy
        self.ax=ax
        self.ay=ay
    def get_x(self):
        return self.x
    def get_y(self):
        return self.y
    def get_vx(self):
        return self.vx
    def get_vy(self):
        return self.vy
    def get_ax(self):
        return self.ax
    def get_ay(self):
        return self.ay

# this fun is used to transfer the ego kinetics state to reference line
# return : map{
# "vx2ref":ego_vx2refline,
# "ax2ref":ego_ax2refline,
# "vy2ref":ego_vy2refline,
# "ay2ref":ego_ay2refline
# }
def transfer2refline(target,ego):
    delta_x=target.get_x()-ego.get_x()
    delta_y=target.get_y()-ego.get_y()
    thirdline=(delta_x**2+delta_y**2)**0.5
    cos_theta=delta_y/thirdline
    sin_theta=delta_x/thirdline
    # lateral velocity and accelerate
    ego_v2refline=ego.get_vx()*cos_theta-ego.get_vy()*sin_theta
    ego_a2refline=ego.get_ax()*cos_theta-ego.get_ay()*sin_theta
    # longitude velocity and accelerate
    ego_vOnrefline=ego.get_vy()*cos_theta+ego.get_vx()*sin_theta
    ego_aOnrefline=ego.get_ay()*cos_theta+ego.get_ax()*sin_theta

    return {"v2ref":ego_v2refline,"a2ref":ego_a2refline,"vOnref":ego_vOnrefline,"aOnref":ego_aOnrefline}

# this fun is used to sample lateral postion
def lateral_sampling(steptime_array,origion_state,d_min,d_max,step):
    d_array=np.arange(d_min,d_max,step)
    # distance, velocity, acceleration
    finalstate_array=[]
    for i in d_array:
        finalstate_array.append([i,0,0])
    
    lateral_coeff_map={}
    for t in steptime_array:
        coeff_list=[]
        for final_state in finalstate_array:
            O=origion_state+final_state
            M=np.array([
                [1,0,0,0,0,0],
                [0,1,0,0,0,0],
                [0,0,2,0,0,0],
                [1,t,t**2,t**3,t**4,t**5],
                [0,1,2*t,3*t**2,4*t**3,5*t**4],
                [0,0,2,6*t,12*t**2,20*t**3]
                ])
            O=np.array([O]).T
            t_co=np.dot(np.linalg.inv(M),O)
            t_co=np.flipud(t_co).flatten()
            t_fun=np.poly1d(t_co.astype(float))
            coeff_list.append(t_fun)
        lateral_coeff_map[t]=coeff_list
    return lateral_coeff_map

# this fun is used to sample longitude velocity
def longitude_sampling(steptime_array,origion_state,v_min,v_max,a_min,a_max,step_num):
    v_array=np.linspace(v_max,v_min,step_num)
    a_array=np.linspace(a_min,a_max,step_num)
    finalstate_array=[]
    for i in range(step_num):
        finalstate_array.append([v_array[i],a_array[i]])

    longitude_coeff_map={}
    for t in steptime_array:
        coeff_list=[]
        for final_longitude_state in finalstate_array:
            O = [0]+origion_state + final_longitude_state
            M=[
                [1,0,0,0,0],    # 初始位移
                [0,1,0,0,0],    # 初始速度
                [0,0,2,0,0],    # 初始加速度
                [0,1, 2*t,3*t**2,4*t**3],
                [0,0,2,6*t,12*t**2]
            ]
            O = np.array([O]).T
            t_co = np.dot(np.linalg.inv(M), O)
            t_co = np.flipud(t_co).flatten()
            t_fun=np.poly1d(t_co.astype(float))
            coeff_list.append(t_fun)
            # store formula not coeff
        longitude_coeff_map[t]=coeff_list
    return longitude_coeff_map

# merge lateral and longitude traj
def merge_traj(lateral_traj_map,longitude_traj_map):
    merged_traj=[]
    for sampletime in lateral_traj_map.keys():
        for lat in lateral_traj_map[sampletime]:
            for long in longitude_traj_map[sampletime]:
                merged_traj.append({"predtime":sampletime,"latTraj":lat,"longTraj":long})
    return merged_traj

# filter invalid lateral and longitude acceleration
def filter_traj(Descartes_traj_set,latAcc_min,latAcc_max,longAcc_min,longAcc_max):
    for index in range(len(Descartes_traj_set)-1,0,-1):
        traj=Descartes_traj_set[index]
        sampleTime=np.arange(0,traj["predtime"],0.02)

        # lat_coeff=traj["latTraj"][::-1]
        # lat_coeff=[x for i in lat_coeff for x in i]
        # diff_lat_coeff=np.polyder(np.polyder(lat_coeff))
        diff_lat=np.polyder(np.polyder(traj["latTraj"]))
        lat_acceleration=np.polyval(diff_lat,sampleTime)
        if (lat_acceleration.max()>latAcc_max or lat_acceleration.min()<latAcc_min):
            Descartes_traj_set.pop(index)
            continue

        long_coeff=traj["longTraj"]
        diff_long=np.polyder(np.polyder(long_coeff))
        long_acceleration=np.polyval(diff_long,sampleTime)
        if (long_acceleration.max()>longAcc_max or long_acceleration.min()<longAcc_min):
            Descartes_traj_set.pop(index)
            # break

# this fun is used to filter trajs which will collide with obj (static obj)
# using SL graphic
def collision_detect(Descartes_traj_set,target_state,ego_state,pred_range):
    # SL Graphic
    smin=sqrt((target_state.get_y()-ego_state.get_y())**2+(target_state.get_x()-ego_state.get_x())**2)
    smax=smin+4.5
    lmin=1
    lmax=-1

    for index in range(len(Descartes_traj_set)-1,-1,-1):
        traj=Descartes_traj_set[index]
        for sampleTime in np.arange(0,pred_range,0.02):
            long_coeff = traj["longTraj"]
            # 积分得到S，三角函数得到d
            s=intg.quad(long_coeff,0,sampleTime)[0]
            lat_coeff = traj["latTraj"]
            d=lat_coeff(sampleTime)

            if (s<smax and s>smin and d<lmin and d>lmax):
                Descartes_traj_set.pop(index)
                break

def optimal_traj(Descartes_traj_set,lat_weight,long_weight):
    min_cost=9999
    optimized_traj={}

    for traj in Descartes_traj_set:
        sampleTime=np.arange(0,traj["predtime"],0.02)

        lat_coeff=traj["latTraj"]
        diff_lat_coeff=np.polyder(np.polyder(np.polyder(lat_coeff)))
        lat_jerk=np.polyval(diff_lat_coeff,sampleTime).sum()

        long_coeff=traj["longTraj"]
        diff_long_coeff=np.polyder(np.polyder(long_coeff))
        long_jerk=np.polyval(diff_long_coeff,sampleTime).sum()

        temp_cost=lat_weight*lat_jerk+long_weight*long_jerk<min_cost
        if (temp_cost<min_cost):
            min=temp_cost
            optimized_traj=traj
    return optimized_traj

def calc_refline(target_state,ego_state):
    target_x=target_state.get_x()
    target_y=target_state.get_y()
    ego_x=ego_state.get_x()
    ego_y=ego_state.get_y()

    k=(target_y-ego_y)/(target_x-ego_x)
    b=target_y-k*target_x
    return {"k":k,"b":b}

def getLocation_in_global(sample_time,traj,refline,ego_state):
    # long_coeff = optimized_traj["longTraj"][::-1]
    # long_coeff = np.array([coeff for i in long_coeff for coeff in i])
    # # 积分得到S，三角函数得到d
    # func = np.poly1d(long_coeff.astype(float))
    long_fun=traj["longTraj"]
    s=intg.quad(long_fun,0,sample_time)[0]

    # lat_coeff = optimized_traj["latTraj"][::-1]
    # lat_coeff = [coeff for i in lat_coeff for coeff in i]
    # d=np.polyval(lat_coeff,sample_time)
    lat_fun=traj["latTraj"]
    L=lat_fun(sample_time)


    sx=math.cos(math.atan(refline["k"]))*s
    sy=math.sin(math.atan(refline["k"]))*s

    sOnref_x=ego_state.get_x()+sx
    sOnref_y=ego_state.get_y()+sy

    dx=math.sin(math.atan(refline["k"]))*L
    dy=math.cos(math.atan(refline["k"]))*L

    global_x=sOnref_x-dx
    global_y=sOnref_y+dy

    return [global_x,global_y]

if __name__ == '__main__':

    # obstacle property
    width=2
    lenght=4.5
    # target location
    target_state=vehicle_state(20.5,60,0,0,0,0)

    # ego original state
    ego_state=vehicle_state(20,0,0,15,0,0.5)

    # trans ego state to reference line
    origion_state = transfer2refline(target_state,ego_state)

    # predict time
    predtime_ar = np.arange(0.5, 5, 0.5)

    refline=calc_refline(target_state,ego_state)
    # lateral planning
    lateral_origion_state=[0,origion_state["v2ref"],origion_state["a2ref"]]
    lateral_coeff_map=lateral_sampling(predtime_ar,lateral_origion_state,0.5,4.5,0.5)
    plt.figure(1)
    for predt in lateral_coeff_map.keys():
        for traj in lateral_coeff_map[predt]:
            x = np.linspace(0, predt, 30)
            y = traj(x)
            plt.plot(x,y)

    # longitude planning
    longitude_origion_state=[origion_state["vOnref"],origion_state["aOnref"]]
    longitude_coeff_map=longitude_sampling(predtime_ar,longitude_origion_state,15,16,0,1,5)
    plt.figure(2)
    for predt in longitude_coeff_map.keys():
        for traj in longitude_coeff_map[predt]:
            x = np.linspace(0, predt, 30)
            y = np.polyder(traj)(x)
            plt.plot(x,y)


    # merge lateral and longitude traj
    merged_traj_list=merge_traj(lateral_coeff_map,longitude_coeff_map)
    print("traj_num = " + str(len(merged_traj_list)))
    filter_traj(merged_traj_list,-2,2,-2,2)
    print("traj_num after filter: " + str(len(merged_traj_list)))
    collision_detect(merged_traj_list,target_state,ego_state,5)
    print("traj_num after collision detect: " + str(len(merged_traj_list)))
    optimized_traj = optimal_traj(merged_traj_list,0.5,0.5)

    plt.figure(3)
    plt.scatter(target_state.get_x(),target_state.get_y())
    plt.scatter(target_state.get_x()-1,target_state.get_y())
    plt.scatter(target_state.get_x()+1,target_state.get_y())
    plt.scatter(ego_state.get_x(),ego_state.get_y())
    for tTraj in merged_traj_list:
        traj_point_x = []
        traj_point_y=[]
        for sampleTime in np.arange(0,tTraj["predtime"],0.2):
            temp_point=getLocation_in_global(sampleTime, tTraj, refline, ego_state)
            traj_point_x.append(temp_point[0])
            traj_point_y.append(temp_point[1])
        plt.plot(traj_point_x,traj_point_y)
    plt.xlim(-10, 50)
    plt.ylim(-10, 200)

    plt.figure(4)
    plt.scatter(target_state.get_x(), target_state.get_y())
    plt.scatter(target_state.get_x() - 1, target_state.get_y())
    plt.scatter(target_state.get_x() + 1, target_state.get_y())
    plt.scatter(ego_state.get_x(), ego_state.get_y())
    traj_point_x = []
    traj_point_y = []
    for sampleTime in np.arange(0, optimized_traj["predtime"], 0.2):
        temp_point = getLocation_in_global(sampleTime, optimized_traj, refline, ego_state)
        traj_point_x.append(temp_point[0])
        traj_point_y.append(temp_point[1])
    plt.plot(traj_point_x, traj_point_y)


    plt.xlim(-10,50)
    plt.ylim(-10,200)
    plt.show()

